#include "PatternCalibrationProcess.h"
#include <glm/gtx/transform.hpp>


PatternCalibrationProcess::PatternCalibrationProcess (
		const ReconstructConfig& config, int checkBoardRows, int checkBoardCols,
		float checkBlockSize) : ProcessBase (config),
		_checkBoard (checkBoardRows, checkBoardCols, checkBlockSize),
		_data (_checkBoard)
{
	_cornersInCamSpace.resize (_checkBoard.getCorners ().size ());
}

PatternCalibrationProcess::~PatternCalibrationProcess ()
{
}

void PatternCalibrationProcess::initialize (SensorResource &sensorRes)
{
	_rp_sensorRes = &sensorRes;
}

void PatternCalibrationProcess::clean ()
{
}

void PatternCalibrationProcess::process ()
{
	if (_recognizePattern (_rp_sensorRes->processedRgb) && !_data.hasInvaildZ ()) {
		_calculateCornersInCameraSpace ();
		glm::mat4x4 rtsmat = _calculateExtrinsicMatrixCv ();
		float focalRatio = 0.0f;
		int acc = 0;
		for (int i = 0; i < _checkBoard.getCornersCount (); ++i) {
			glm::vec4 standardCornerInCamSpace = rtsmat * glm::vec4 (
					_checkBoard.getCorners ()[i].x,
					_checkBoard.getCorners ()[i].y,
					_checkBoard.getCorners ()[i].z, 1.0f);
			if (_data.zAccArr[i] != 0) {
				focalRatio += _cornersInCamSpace[i].z / standardCornerInCamSpace.z;
				++acc;
			}
		}
		if (acc != 0) {
			focalRatio /= (float)acc;
			SensorIntrinsic newIntrinsic = _rp_sensorRes->rp_sensor->getIntrinsic ();
			glm::mat3 intMat = newIntrinsic.getMatrix ();
			intMat[0][0] *= focalRatio;
			intMat[1][1] *= focalRatio;
			newIntrinsic.set (intMat, newIntrinsic.getDistCeoffs ());
			_rp_sensorRes->rp_sensor->setIntrinsic (newIntrinsic);
		}
		_rp_sensorRes->rp_sensor->setExtrinsicMat (rtsmat);
		_showResult (_rp_sensorRes->processedRgb, rtsmat);
	}

}

bool PatternCalibrationProcess::_recognizePattern (const QImage& image)
{
	if (image.isNull ())
		return false;

	
	bool isPattFound;
	std::vector< cv::Point2f > corners = _checkBoard.findCheckBoard (
			image, isPattFound);
	if (isPattFound) {
		for (int i = 0; i < corners.size () / 2; ++i) {
			cv::Point2f temp = corners[i];
			corners[i] = corners[corners.size () - 1 - i];
			corners[corners.size() - 1 - i] = temp;
		}
		for (int i = 0; i < _checkBoard.getRows () * _checkBoard.getCols (); ++i) {
			if (cv::norm (_data.cornersAcc[i] - corners[i]) >= 4.0f) {
				_data.cornersAcc[i] = corners[i];
				_cornersInCamSpace[i] = _rp_sensorRes->rp_sensor->get3DPointInCam (
						_data.cornersAcc[i].x, _data.cornersAcc[i].y);
				_data.accArr[i] = 1;
				if (_cornersInCamSpace[i].z < 100.0f)
					_data.zAccArr[i] = 0;
				else
					_data.zAccArr[i] = 1;
			} else {
				++_data.accArr[i];
				float srcPercent = 1.0f / (float)_data.accArr[i];
				float dstPercent = 1.0f - srcPercent;
				_data.cornersAcc[i] = _data.cornersAcc[i] * dstPercent +
						corners[i] * srcPercent;
				glm::vec3 newPoint = _rp_sensorRes->rp_sensor->get3DPointInCam (
						_data.cornersAcc[i].x, _data.cornersAcc[i].y);
				if (newPoint.z >= 100.0f) {
					if (_data.zAccArr[i] == 0) {
						_cornersInCamSpace[i] = newPoint;
						_data.zAccArr[i] = 1;
					} else {
						++_data.zAccArr[i];
						float srcPercent = 1.0f / (float)_data.zAccArr[i];
						float dstPercent = 1.0f - srcPercent;
						_cornersInCamSpace[i] = 
								_cornersInCamSpace[i] * dstPercent +
								newPoint * srcPercent;
					}
				}
			}
		}
	}
	return isPattFound;
}

void PatternCalibrationProcess::_calculateCornersInCameraSpace ()
{
	// glm::vec3 corners3D[_checkBoard.getCornersCount ()];
}

void PatternCalibrationProcess::_showResult (const QImage& image, glm::mat4x4 rtsmat)
{
	glm::mat3 perspective = _rp_sensorRes->rp_sensor->getIntrinsic ().getMatrix ();

	glm::mat4x4 rtsmatR = glm::inverse (rtsmat);
	glm::mat4x3 rtsmat34 (rtsmat[0].swizzle (glm::X, glm::Y, glm::Z),
			rtsmat[1].swizzle (glm::X, glm::Y, glm::Z),
			rtsmat[2].swizzle (glm::X, glm::Y, glm::Z),
			rtsmat[3].swizzle (glm::X, glm::Y, glm::Z));
	glm::vec3* result = new glm::vec3[_checkBoard.getCornersCount ()];
	glm::vec4* objPointsInCamera = new glm::vec4[_checkBoard.getCornersCount ()];
	glm::vec4* cornersInWorld = new glm::vec4[_checkBoard.getCornersCount ()];
	std::vector< cv::Point2f > result2D (_checkBoard.getCornersCount ());

	for (int i = 0; i < _checkBoard.getCornersCount (); ++i) {
		objPointsInCamera[i] = rtsmat * glm::vec4 (_checkBoard.getCorners ()[i].x,
				_checkBoard.getCorners ()[i].y, _checkBoard.getCorners ()[i].z, 1);
		cornersInWorld[i] = rtsmatR * glm::vec4 (_cornersInCamSpace[i], 1.0);
		result[i] = perspective * (rtsmat34 * glm::vec4 (
				_checkBoard.getCorners ()[i].x, _checkBoard.getCorners ()[i].y,
				_checkBoard.getCorners ()[i].z, 1));
		result[i] /= result[i].z;
		result2D[i] = cv::Point2f (result[i].x, result[i].y);
	}
	static cv::Mat _pattImageShow (image.height (), image.width (), CV_8UC4);
	memcpy (_pattImageShow.data, image.bits (), image.byteCount ());
	//cvCopy (image, _pattImageShow);
	cv::drawChessboardCorners (_pattImageShow,
			cv::Size (_checkBoard.getCols (), _checkBoard.getRows ()), _data.cornersAcc, true);
	imshow ((QString ("Pattern Image - ") + _rp_sensorRes->rp_sensor->getName ()).
			toLocal8Bit ().data (), _pattImageShow);

	memcpy (_pattImageShow.data, image.bits (), image.byteCount ());
	//cvCopy (image, _pattImageShow);
	cv::drawChessboardCorners (_pattImageShow,
			cv::Size (_checkBoard.getCols (), _checkBoard.getRows ()), result2D, true);
	imshow ((QString ("New Pattern Image - ") + _rp_sensorRes->rp_sensor->getName ()).
			toLocal8Bit ().data (), _pattImageShow);

	delete [] result;
	delete [] objPointsInCamera;
	delete [] cornersInWorld;
}

glm::mat4 PatternCalibrationProcess::_calculateExtrinsicMatrix ()
{
	return _calculateTranslation () * _calculateRotation () * _calculateScale ();
}


glm::mat4 PatternCalibrationProcess::_calculateTranslation ()
{
	return glm::translate (_cornersInCamSpace[(_checkBoard.getRows () - 1) * _checkBoard.getCols ()]);
}

glm::mat4 PatternCalibrationProcess::_calculateRotation ()
{
	glm::vec3 xVec;
	glm::vec3 yVec;
	int xCount = 0;
	int yCount = 0;

	for (int row = 0; row < _checkBoard.getRows (); ++row) {
		for (int col1 = 0; col1 < _checkBoard.getCols (); ++col1) {
			glm::vec3 p1 = _cornersInCamSpace[row * _checkBoard.getCols () + col1];
			for (int col2 = col1 + 3; col2 < _checkBoard.getCols (); ++col2) {
				glm::vec3 p2 = _cornersInCamSpace[row * _checkBoard.getCols () + col2];
				glm::vec3 v12 = (p2 - p1) / (float)((col2 - col1) * _checkBoard.getBlockSize ());
				xVec = (xVec * (float)xCount + v12) / (float)(xCount + 1);
				++xCount;
			}
		}
	}

	for (int col = 0; col < _checkBoard.getCols (); ++col) {
		for (int row1 = 0; row1 < _checkBoard.getRows (); ++row1) {
			glm::vec3 p1 = _cornersInCamSpace[row1 * _checkBoard.getCols () + col];
			for (int row2 = row1 + 3; row2 < _checkBoard.getRows (); ++row2) {
				glm::vec3 p2 = _cornersInCamSpace[row2 * _checkBoard.getCols () + col];
				glm::vec3 v21 = (p1 - p2) / (float)((row2 - row1) * _checkBoard.getBlockSize ());
				yVec = (yVec * (float)yCount + v21) / (float)(yCount + 1);
				++yCount;
			}
		}
	}

	glm::vec3 zVec = -glm::normalize (glm::cross (xVec, yVec));
	glm::mat4 rmat (glm::vec4 (xVec, 0.0), glm::vec4 (yVec, 0.0),
			glm::vec4 (zVec, 0.0), glm::vec4 (0.0, 0.0, 0.0, 1.0));
	return rmat;
}

glm::mat4 PatternCalibrationProcess::_calculateScale ()
{
	float _scale = 0.0;
	std::vector< float > scaleArr;
	float minScale = 100;
	float maxScale = 0;
	int count = 0;
	for (int i = 0; i < _checkBoard.getCornersCount (); ++i) {
		for (int j = i + 1; j < _checkBoard.getCornersCount (); ++j) {
			float distanceWorld = cv::norm (_checkBoard.getCorners ()[i] -
					_checkBoard.getCorners ()[j]);
			if (distanceWorld < 4) continue;
			float distanceCamera = glm::distance (_cornersInCamSpace[i],
					_cornersInCamSpace[j]);
			float curScale = distanceCamera / distanceWorld;

			_scale = (_scale * count + curScale) / (count + 1);
			++count;

			scaleArr.push_back (curScale);
			if (curScale < minScale)
				minScale = curScale;
			if (curScale > maxScale)
				maxScale = curScale;
		}
	}
	_data.scale = _scale;
	return glm::scale (1.0f, 1.0f, _scale);
}

glm::mat4 PatternCalibrationProcess::_calculateExtrinsicMatrixCv ()
{
	CvMat objectPoints = cvMat (_checkBoard.getCornersCount (), 1, CV_32FC3,
			(void*)&_checkBoard.getCorners ()[0].x);
	CvMat imagePoints = cvMat (_checkBoard.getCornersCount (), 1, CV_32FC2,
			&_data.cornersAcc[0].x);
	glm::vec3 rVec, tVec;
	glm::mat3 rotateMatrix;
	CvMat rVecM = cvMat (3, 1, CV_32FC1, &rVec[0]);
	CvMat rMat = cvMat (3, 3, CV_32FC1, &rotateMatrix[0][0]);
	CvMat tVecM = cvMat (3, 1, CV_32FC1, &tVec[0]);
	CvMat cameraMat = cvMat (3, 3, CV_32FC1,
			(void*)&_rp_sensorRes->rp_sensor->getIntrinsic ().matrixCvMat ()(0, 0));
	cvFindExtrinsicCameraParams2 (&objectPoints, &imagePoints,
			&cameraMat, NULL, &rVecM, &tVecM);
	cvRodrigues2 (&rVecM, &rMat);
 	_calculateScale ();
	glm::mat4 rtsMatrix (glm::transpose (rotateMatrix));
	rtsMatrix = glm::translate (tVec) * rtsMatrix;
	rtsMatrix = glm::scale (_data.scale, _data.scale, _data.scale) * rtsMatrix;
	return rtsMatrix;
}

glm::mat4 PatternCalibrationProcess::_calculateExtrinsicMatrixBySolveEquation ()
{
	for (int i = 0, index = _checkBoard.getCornersCount (); i < std::min (_checkBoard.getRows (), _checkBoard.getCols ()); ++i) {
		glm::vec3 startPoint = _cornersInCamSpace[i * _checkBoard.getCols () + i];
		for (int j = 1; j < std::min (_checkBoard.getRows (), _checkBoard.getCols ()) - i; ++j, ++index) {
			glm::vec3 p1 = _cornersInCamSpace[i * _checkBoard.getCols () + i + j];
			glm::vec3 v1 = p1 - startPoint;
			glm::vec3 p2 = _cornersInCamSpace[(i + j) * _checkBoard.getCols () + i];
			glm::vec3 v2 = p2 - startPoint;
			glm::vec3 vz = glm::cross (v1, v2);
			vz = glm::normalize (vz) * _checkBoard.getBlockSize () * (float)j;
			_cornersInCamSpace[index] = startPoint + vz;
		}
	}

	CvMat *pCornersInCamH = cvCreateMat (_checkBoard.getCornersCount () + _checkBoard.getDepthSegs (), 4, CV_32FC1);
	for (int i = 0; i < _checkBoard.getCornersCount () + _checkBoard.getDepthSegs (); ++i) {
		CV_MAT_ELEM (*pCornersInCamH, float, i, 0) = _cornersInCamSpace[i].x;
		CV_MAT_ELEM (*pCornersInCamH, float, i, 1) = _cornersInCamSpace[i].y;
		CV_MAT_ELEM (*pCornersInCamH, float, i, 2) = _cornersInCamSpace[i].z;
		CV_MAT_ELEM (*pCornersInCamH, float, i, 3) = 1.0f;
	}
	CvMat cornersInWorld = cvMat (_checkBoard.getCornersCount () + _checkBoard.getDepthSegs (), 3, CV_32FC1,
			(void*)&_checkBoard.getCorners ()[0].x);
	CvMat *pTransformMat = cvCreateMat (4, 3, CV_32FC1);
	cvSolve (pCornersInCamH, &cornersInWorld, pTransformMat, CV_SVD);

	glm::mat4 rtsMatrix;
	for (int col = 0; col < 4; ++col)
		for (int row = 0; row < 3; ++row)
			rtsMatrix[col][row] = CV_MAT_ELEM (*pTransformMat, float, col, row);
	rtsMatrix = glm::inverse (rtsMatrix);
	cvReleaseMat (&pCornersInCamH);
	cvReleaseMat (&pTransformMat);
	return rtsMatrix;
}

